package uestc.lj.basic.graph;

import java.util.*;

/**
 * 最小生成树算法
 *
 * @Author:Crazlee
 * @Date:2021/11/24
 */
public class Code04_Kruskal {

	/**
	 * 创建每个节点的集合
	 * 优化：利用并查集
	 */
	public static class MySets {
		public HashMap<Node, List<Node>> setMap;

		public MySets(List<Node> nodes) {
			for (Node node : nodes) {
				List<Node> set = new ArrayList<>();
				set.add(node);
				setMap.put(node, set);
			}
		}

		/**
		 * 判断from节点和to节点是否在一个集合中
		 *
		 * @param from 节点
		 * @param to   节点
		 * @return 是否存在环
		 */
		public boolean isSameSet(Node from, Node to) {
			List<Node> fromSet = setMap.get(from);
			List<Node> toSet = setMap.get(to);
			//如果内存地址一样则一定是一个集合
			return fromSet == toSet;
		}

		/**
		 * 合并两个节点的集合
		 *
		 * @param from
		 * @param to
		 */
		public void union(Node from, Node to) {
			List<Node> fromSet = setMap.get(from);
			List<Node> toSet = setMap.get(to);
			for (Node node : toSet) {
				fromSet.add(node);
				setMap.put(node, fromSet);
			}
		}
	}

	/**
	 * 克鲁斯卡尔算法生成最小生成树
	 *
	 * @param graph
	 * @return
	 */
	public static Set<Edge> kruskalMST(Graph graph) {
		List<Node> nodes = new ArrayList<>();
		for (Node node : graph.nodes.values()) {
			nodes.add(node);
		}
		MySets mySets = new MySets(nodes);
		PriorityQueue<Edge> priorityQueue = new PriorityQueue<>(new Comparator<Edge>() {
			@Override
			public int compare(Edge o1, Edge o2) {
				return o1.weight - o2.weight;
			}
		});
		priorityQueue.addAll(graph.edges);
		Set<Edge> result = new HashSet<>();
		while (!priorityQueue.isEmpty()) {
			Edge edge = priorityQueue.poll();
			if (!mySets.isSameSet(edge.from, edge.to)) {
				result.add(edge);
				mySets.union(edge.from, edge.to);
			}
		}
		return result;
	}

	//优化使用并查集
	public static class UnionFind {
		private HashMap<Node, Node> fatherMap;
		private HashMap<Node, Integer> rankMap;

		public UnionFind() {
			fatherMap = new HashMap<Node, Node>();
			rankMap = new HashMap<Node, Integer>();
		}

		private Node findFather(Node n) {
			Node father = fatherMap.get(n);
			if (father != n) {
				father = findFather(father);
			}
			fatherMap.put(n, father);
			return father;
		}

		public void makeSets(Collection<Node> nodes) {
			fatherMap.clear();
			rankMap.clear();
			for (Node node : nodes) {
				fatherMap.put(node, node);
				rankMap.put(node, 1);
			}
		}

		public boolean isSameSet(Node a, Node b) {
			return findFather(a) == findFather(b);
		}

		public void union(Node a, Node b) {
			if (a == null || b == null) {
				return;
			}
			Node aFather = findFather(a);
			Node bFather = findFather(b);
			if (aFather != bFather) {
				int aFrank = rankMap.get(aFather);
				int bFrank = rankMap.get(bFather);
				if (aFrank <= bFrank) {
					fatherMap.put(aFather, bFather);
					rankMap.put(bFather, aFrank + bFrank);
				} else {
					fatherMap.put(bFather, aFather);
					rankMap.put(aFather, aFrank + bFrank);
				}
			}
		}
	}

	public static class EdgeComparator implements Comparator<Edge> {

		@Override
		public int compare(Edge o1, Edge o2) {
			return o1.weight - o2.weight;
		}

	}

	public static Set<Edge> kruskalMSTbyUnionFind(Graph graph) {
		UnionFind unionFind = new UnionFind();
		unionFind.makeSets(graph.nodes.values());
		PriorityQueue<Edge> priorityQueue = new PriorityQueue<>(new EdgeComparator());
		for (Edge edge : graph.edges) {
			priorityQueue.add(edge);
		}
		Set<Edge> result = new HashSet<>();
		while (!priorityQueue.isEmpty()) {
			Edge edge = priorityQueue.poll();
			if (!unionFind.isSameSet(edge.from, edge.to)) {
				result.add(edge);
				unionFind.union(edge.from, edge.to);
			}
		}
		return result;
	}
}

